# -*- coding: utf-8 -*-
import json,time,sqlite3
from datetime import datetime,timedelta
import threading
from threading import Thread
from pysnmp.entity.rfc3413.oneliner import cmdgen
import traceback

from Queue import Queue
#具体的处理函数，负责处理单个任务
def do_somthing_using(arguments):
    print 'starting:',arguments['ndip']
    global nowflow
    flow={}
    flow['id']=arguments['id']
    flow['ircv'],flow['isnt'],flow['icpe'],flow['dtlog']=get_flow(arguments['ndip'],arguments['iport'])
    nowflow.append(flow)




#这个是工作进程，负责不断从队列取数据并处理
def working():
    while True:
        arguments = q.get()
        try:
            do_somthing_using(arguments)
        except:
            pass
        q.task_done()
def ndflow_pool(nds):
    s=time.time()
    #把JOBS排入队列
    for i in nds:
        q.put(i)
        #等待所有JOBS完成
    q.join()
    print 'pool time',time.time()-s

def getsnmp(ndip,oid):
    try:
        errorIndication, errorStatus, errorIndex, varBinds = cmdgen.CommandGenerator().getCmd(cmdgen.CommunityData(ndip, 'public'),cmdgen.UdpTransportTarget((ndip, 161)),oid,)
        if errorIndication:
            print errorIndication
            return None
        else:
            if errorStatus:
                print '%s at %s\n' % (
                    errorStatus.prettyPrint(),
                    errorIndex and varBinds[int(errorIndex)-1] or '?'
                    )
                return None
            else:
                for name, val in varBinds:
                    # print '%s:%s = %s' % (ndip,name.prettyPrint(), val.prettyPrint())
                    # snmp=val.prettyPrint()
                    return val.prettyPrint()
    except:
        return None

def get_flow(ndip,iport):
    roid=[1,3,6,1,4,1,5959,1,1,2,1,1,3,1]
    roid.append(iport)
    roid=tuple(roid)
    soid=[1,3,6,1,4,1,5959,1,1,2,1,1,8,1]
    soid.append(iport)
    soid=tuple(soid)
    coid=[1,3,6,1,4,1,5959,1,3,1,1,1,2,1]
    # coid.append(iport)
    coid=tuple(coid)
    rcv=getsnmp(ndip,roid)
    snt=getsnmp(ndip,soid)
    cpe=getsnmp(ndip,coid)
    log=datetime.strftime(datetime.now(),'%Y-%m-%d %H:%M:%S.%f')
    return rcv,snt,cpe,log

def find_nd(nds,id):
    if type(nds)==list:
        for nd in nds:
            if nd['id']==id:
                return nd

def calc_flow(ni,pi,deltatime):
    if deltatime<=0:
        return 0
    if ni<pi:
        if ni+pi>4294967296:
#            计数器是32位整数，溢出后重新开始计数，因此要计算准确的值
            x=4294967296-pi+ni
        else:
#            重启后计数器清零，所以以现在的字节数为准
            x=ni
    else:
        x=ni-pi
    return x*8/deltatime

def calc_cpe(ni,pi):
    if ni is None:
        if pi is None:
            return 0
        else:
            return pi
    else:
        return ni

def save_flows(prevflow,nowflow):
    s=time.time()
    while True:
        try:
            flows=[]
#            print prevflow,nowflow
            for n in nowflow:
                if (n['ircv'] is not None) and (n['isnt'] is not None):
#                    print n
                    p=find_nd(prevflow,n['id'])
                    # print 'p',p
                    # print 'N',n
                    frx=ftx=0
                    icpe=calc_cpe(n['icpe'],p['icpe'])
                    if p['dtlog']:
                        plog=datetime.strptime(p['dtlog'],'%Y-%m-%d %H:%M:%S.%f')
                        nlog=datetime.strptime(n['dtlog'],'%Y-%m-%d %H:%M:%S.%f')
                        deltatime=timedelta.total_seconds(nlog-plog)
                        nircv=int(n['ircv'])
                        nisnt=int(n['isnt'])
                        pircv=int(p['ircv'])
                        pisnt=int(p['isnt'])
                        frx=calc_flow(nircv,pircv,deltatime)
                        ftx=calc_flow(nisnt,pisnt,deltatime)
                    flows.append((n['id'],n['dtlog'],n['ircv'],n['isnt'],frx,ftx,icpe))
            con=sqlite3.connect(r'ndflow.db3')
            con.row_factory = sqlite3.Row
            cur=con.cursor()
#            print flows
            for flow in flows:
                # print flow
                cur.execute('insert into ndflowlog(ndflowid,dtlog,ircv,isnt,frx,ftx,icpe) values(?,?,?,?,?,?,?)',flow)
            con.commit()
            cur.close()
            con.close()
            print 'write db',time.time()-s
            return
        except :
            print traceback.format_exc()
            time.sleep(1)

def get_prevflow():
    s=time.time()
    while True:
        try:
            prevndflows=[]
            con=sqlite3.connect(r'ndflow.db3')
            con.row_factory = sqlite3.Row
            cur=con.cursor()
            now=datetime.strftime(datetime.now(),'%Y-%m-%d %H:%M:%S.%f')
            nds=cur.execute("select id,ndip,iport,dtend from ndflow where dtstart<=? and ?<=dtend",(now,now)).fetchall()
            for nd in nds:
                ndflow={}
                ndflow['id']=nd['id']
                ndflow['ndip']=nd['ndip']
                ndflow['iport']=nd['iport']
                ndflow['dtend']=nd['dtend']
                prev=cur.execute('select ircv,isnt,dtlog,icpe from ndflowlog where ndflowid=? order by id desc limit 1',(nd['id'],)).fetchone()
                if prev:
                    ndflow['ircv']=prev['ircv']
                    ndflow['isnt']=prev['isnt']
                    ndflow['dtlog']=prev['dtlog']
                    ndflow['icpe']=prev['icpe']
                else:
                    ndflow['ircv']=None
                    ndflow['isnt']=None
                    ndflow['dtlog']=None
                    ndflow['icpe']=None
                prevndflows.append(ndflow)
            cur.close()
            con.close()
            print 'read db',time.time()-s
            return prevndflows
        except:
            print traceback.format_exc()
            time.sleep(1)


q = Queue()
#fork NUM个线程等待队列
NUM = 10
for i in range(NUM):
    t = Thread(target=working)
    t.setDaemon(True)
    t.start()



while True:
    config = json.load(open("config.txt",'r'))
    delaytime=config["NDFlow"]
    starttime=time.time()
    prevflow=get_prevflow()
    nowflow=[]

    ndflow_pool(prevflow)
    print "threading:",threading.activeCount()

    save_flows(prevflow,nowflow)
    endtime=time.time()
    sleeptime=delaytime-(endtime-starttime)
    print "sleep ndflow",sleeptime
    if sleeptime>0:
        time.sleep(sleeptime)
